/**
 * @file ebex_genetic.c
 *
 *
 * @date Jan 12, 2010
 * @author seth
 *
 * @brief This file is part of EBEX State Estimator, created for the EBEX project
 *
 * This software is copyright (C) 2010 Columbia University
 *
 * EBEX State Estimator is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * EBEX State Estimator is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with EBEX State Estimator; if not, write to the Free Software Foundation, Inc.,
 * 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */

#include <ebex_sort.h>
#include <ebex_genetic_chromosome.h>
#include <ebex_genetic.h>
#include <ebex_genetic_stats.h>

static bool ebex_filter_openmp_initialised = false;


population_t *genetic_pop_new(const int stable_size, const int num_chromosome)
{
	population_t *newpop = NULL; /* New population structure. */
	int i; /* Loop over (unassigned) entities. */

	if (!(newpop = e_malloc(sizeof(population_t))))
		die("Unable to allocate memory");

	e_memset(newpop, 0, sizeof(population_t));

	newpop->stable_size = stable_size;
	newpop->max_size = (1 + stable_size) * 4; /* +1 prevents problems if stable_size is 0. */

	newpop->num_chromosomes = num_chromosome;

	newpop->data = NULL;
	newpop->free_index = newpop->max_size - 1;

	newpop->crossover_ratio = GENETIC_DEFAULT_CROSSOVER_RATIO;
	newpop->mutation_ratio = GENETIC_DEFAULT_MUTATION_RATIO;
	newpop->scheme = EVOLUTION_MODE_DARWIN;
	newpop->elitism = GENETIC_SELECTION_PARENTS_SURVIVE;


	omp_init_lock(&newpop->lock);
	omp_init_lock(&newpop->chromo_chunk_lock);

	if (!(newpop->raw_organisms = e_malloc(newpop->max_size*sizeof(organism_t*))))
		die("Unable to allocate memory");

	if (!(newpop->sorted_organisms = e_malloc(newpop->max_size*sizeof(organism_t*))))
		die("Unable to allocate memory");

	newpop->organism_memory = mem_chunk_new(sizeof(organism_t), 512);

	/*
	 * Wipe the entity arrays.
	 */
	for (i = 0; i < newpop->max_size; i++)
	{
		newpop->raw_organisms[i] = NULL;
		newpop->sorted_organisms[i] = NULL;
	}

	return newpop;
}

bool genetic_organism_seed(population_t *m_pop, organism_t *adam)
{

	if (!m_pop)
		die("Null pointer to population structure passed.");
	if (!m_pop->seed)
		die("Population seeding function is not defined.");

	return m_pop->seed(m_pop, adam);
}

bool genetic_fill_pop(population_t *m_pop, int num)
{
	int i; /* Loop variables. */
	organism_t *adam; /* New population member. */

	if (!m_pop)
		die("Null pointer to population structure passed.");
	if (!m_pop->seed)
		die("Population seeding function is not defined.");

#	pragma omp parallel for \
		shared(m_pop, num) private(i,adam) schedule(static)
	for (i = 0; i < num; i++)
	{
		adam = genetic_get_new_organism(m_pop);
		m_pop->seed(m_pop, adam);
	}

	return true;
}


double genetic_organism_evaluate(population_t *m_pop, organism_t *this_entity)
{

	if (!m_pop)
		die("Null pointer to population structure passed.");
	if (!this_entity)
		die("Null pointer to entity structure passed.");
	if (!m_pop->evaluate)
		die("Evaluation callback not defined.");

	if (m_pop->evaluate(m_pop, this_entity) == false
		)
		this_entity->fitness = GENETIC_MIN_FITNESS;

	return this_entity->fitness;
}

int genetic_organism_get_rank(population_t *m_pop, organism_t *e)
{
	size_t rank = 0;
	do
	{
		if (m_pop->sorted_organisms[rank] == e) return rank;
	}while (++rank < m_pop->size);

	return -1;
}

int genetic_get_index(population_t *m_pop, organism_t *e)
{
	size_t id = 0;

	if (!m_pop)
		die("Null pointer to population structure passed.");
	if (!e)
		die("Null pointer to entity structure passed.");

	while (id < m_pop->max_size)
	{
		if (m_pop->raw_organisms[id] == e)
			return id;
		id++;
	}

	return -1;
}

organism_t *genetic_get_organism(population_t *m_pop, const size_t id)
{
	if (!m_pop)
		die("Null pointer to population structure passed.");

	if (id > m_pop->max_size)
		return NULL;

	return m_pop->raw_organisms[id];
}


organism_t *genetic_get_ranked_organism(population_t *m_pop, const size_t rank)
{
	if (!m_pop)
		die("Null pointer to population structure passed.");

	if (rank > m_pop->size)
		return NULL;

	return m_pop->sorted_organisms[rank];
}

static bool genetic_init_organism(population_t *m_pop, organism_t *m_new)
{

	if (!m_new)
		die("Null pointer to entity structure passed.");
	if (!m_pop->chromosome_constructor)
		die("Chromosome constructor not defined.");

	/* Allocate chromosome structures. */
	m_new->chromosome = NULL;
	m_new->allele_data = NULL;
	m_new->allele_size = 0;
	m_new->allocated_chromosomes = 0;
	m_pop->chromosome_constructor(m_pop, m_new);

	/* Physical characteristics currently undefined. */
	m_new->data = NULL;

	/* No fitness evaluated yet. */
	m_new->fitness = GENETIC_MIN_FITNESS;
	m_new->previous_weight = 1.0;

	return true;
}

bool genetic_free_rank(population_t *m_pop, int m_rank)
{

	organism_t *corpse = m_pop->sorted_organisms[m_rank]; /* Dead entity. */

	if (!corpse)
		die("Invalid entity rank");

	/* Clear user data. */
	if (corpse->data)
	{
		if(m_pop->organism_data_destructor) m_pop->organism_data_destructor(corpse->data);
		corpse->data = NULL;
	}

	omp_set_lock(&m_pop->lock);

	m_pop->size--;
	if (corpse->chromosome)
	{
		m_pop->chromosome_destructor(m_pop, corpse);
		corpse->chromosome = NULL;
	}

	memmove(&m_pop->sorted_organisms[m_rank], &m_pop->sorted_organisms[m_rank+1], (m_pop->size - m_rank) * sizeof(organism_t*));

	m_pop->sorted_organisms[m_pop->size] = NULL;

	m_pop->raw_organisms[genetic_get_index(m_pop, corpse)] = NULL;

	omp_unset_lock(&m_pop->lock);

	mem_chunk_free(m_pop->organism_memory, corpse);

	return true;
}

bool genetic_free_index(population_t *m_pop, int m_id)
{
	size_t i;
	organism_t *corpse = m_pop->raw_organisms[m_id];

	if (!corpse)
		die("Invalid entity index");

	if (corpse->data)
	{
		if(m_pop->organism_data_destructor) m_pop->organism_data_destructor(corpse->data);
		corpse->data = NULL;
	}

	omp_set_lock(&m_pop->lock);

	/* Population size is one less now! */
	m_pop->size--;

	/* Update entity_iarray[], so there are no gaps! */
	for (i = genetic_organism_get_rank(m_pop, corpse); i < m_pop->size; i++)
		m_pop->sorted_organisms[i] = m_pop->sorted_organisms[i + 1];

	m_pop->sorted_organisms[m_pop->size] = NULL;

	if (corpse->chromosome)
		m_pop->chromosome_destructor(m_pop, corpse);

	omp_unset_lock(&m_pop->lock);

	m_pop->raw_organisms[m_id] = NULL;
	mem_chunk_free(m_pop->organism_memory, corpse);

	return true;
}

bool genetic_free_organism(population_t *m_pop, organism_t *dying)
{
	return genetic_free_rank(m_pop, genetic_organism_get_rank(m_pop, dying));
}

void genetic_clear_organism(population_t *p, organism_t *this_entity)
{
	if (this_entity->data)
	{
		p->organism_data_destructor(this_entity);

		this_entity->data = NULL;
	}

	this_entity->fitness = GENETIC_MIN_FITNESS;
	return;
}

/**
 * Gets an unused organism pointer from our memory pool.  If none available, allocate additional entities.
 * @param m_pop Population to use for the organism.
 * @return Pointer to the fresh organism
 */
organism_t *genetic_get_new_organism(population_t *m_pop)
{
	int new_max_size; /* Increased maximum number of entities. */
	int i;
	organism_t *fresh; /* Unused entity structure. */

	omp_set_lock(&m_pop->lock);

	if (m_pop->max_size == (m_pop->size + 1))
	{ /* No, so allocate some more space. */
		elog(log_level_verbose, "No unused entities available -- allocating additional structures.");

		new_max_size = (m_pop->max_size * 3) / 2 + 1;
		m_pop->raw_organisms = e_realloc(m_pop->raw_organisms, new_max_size*sizeof(organism_t*));
		m_pop->sorted_organisms = e_realloc(m_pop->sorted_organisms, new_max_size*sizeof(organism_t*));

		for (i = m_pop->max_size; i < new_max_size; i++)
		{
			m_pop->raw_organisms[i] = NULL;
			m_pop->sorted_organisms[i] = NULL;
		}

		m_pop->max_size = new_max_size;
		m_pop->free_index = new_max_size - 1;
	}

	/* Find unused entity index. */
	while (m_pop->raw_organisms[m_pop->free_index])
	{
		if (m_pop->free_index == 0)
			m_pop->free_index = m_pop->max_size;
		m_pop->free_index--;
	}

	/* Prepare it. */
	fresh = (organism_t *) mem_chunk_alloc(m_pop->organism_memory);

	m_pop->raw_organisms[m_pop->free_index] = fresh;
	genetic_init_organism(m_pop, fresh);

	/* Store in lowest free slot in entity_iarray */
	m_pop->sorted_organisms[m_pop->size] = fresh;

	/* Population is bigger now! */
	m_pop->size++;

	omp_unset_lock(&m_pop->lock);

	return fresh;
}


bool ebex_genetic_copy_data(population_t *m_pop, organism_t *dest, organism_t *src, const int chromosome)
{
	if (m_pop->chromosome_data_copy) dest->chromosome[chromosome].user_data = m_pop->chromosome_data_copy(src->chromosome[chromosome].user_data);
	else dest->chromosome[chromosome].user_data = src->chromosome[chromosome].user_data;

	return true;
}

size_t ebex_genetic_copy_chromosome(population_t *m_pop, organism_t *dest, organism_t *src, const size_t m_chromosome)
{
	ebex_genetic_copy_data(m_pop, dest, src, m_chromosome);
	dest->chromosome[m_chromosome].allele_max_double = src->chromosome[m_chromosome].allele_max_double;
	dest->chromosome[m_chromosome].allele_min_double = src->chromosome[m_chromosome].allele_min_double;
	dest->chromosome[m_chromosome].allele_mutation_prob = src->chromosome[m_chromosome].allele_mutation_prob;
	dest->chromosome[m_chromosome].length = src->chromosome[m_chromosome].length;
	dest->chromosome[m_chromosome].type = src->chromosome[m_chromosome].type;
	switch(src->chromosome[m_chromosome].type)
	{
		case CHROMOSOME_TYPE_DOUBLE:
			return ebex_genetic_chromosome_double_copy(m_pop, src, dest, m_chromosome);
			break;
		case CHROMOSOME_TYPE_QUATERNION:
			return ebex_genetic_chromosome_quaternion_copy(m_pop, src, dest, m_chromosome);
			break;
		default:
			break;
	}
	return 0;
}

bool genetic_organism_copy_all_chromosomes(population_t *m_pop, organism_t *dest, organism_t *src)
{
	size_t i;

	if (!dest || !src)
		die("Null pointer to entity structure passed");

	for (i = 0; i < m_pop->num_chromosomes; i++)
	{
		ebex_genetic_copy_chromosome(m_pop, dest, src, i);
	}

	return true;
}

bool genetic_organism_copy(population_t *m_pop, organism_t *dest, organism_t *src)
{

	genetic_organism_copy_all_chromosomes(m_pop, dest, src);
	dest->fitness = src->fitness;
	dest->previous_weight = src->previous_weight;

	return true;
}

organism_t *genetic_organism_clone(population_t *m_pop, organism_t *parent)
{
	organism_t *dolly; /* The clone. */

	dolly = genetic_get_new_organism(m_pop);

	genetic_organism_copy(m_pop, dolly, parent);

	return dolly;
}


void genetic_pop_set_parameters(population_t *m_pop, const evolution_mode_t scheme, const genetic_selection_mode_t elitism,
		const double crossover, const double mutation, const double migration)
{

	if (!m_pop)
		die("Null pointer to population structure passed.");

	elog(
			log_level_verbose,
			"Population's parameters: scheme = %d elitism = %d crossover = %f mutation = %f migration = %f", (int) scheme, (int) elitism, crossover, mutation, migration);

	m_pop->scheme = scheme;
	m_pop->elitism = elitism;
	m_pop->crossover_ratio = crossover;
	m_pop->mutation_ratio = mutation;
}

void genetic_organism_set_allele_mutation_prob(organism_t *m_entity, const int m_num_chrom, const double m_prob)
{

	if (!m_entity)
		die("Null pointer to entity structure passed.");

	for (int i = 0; i < m_num_chrom; i++)
	{
		m_entity->chromosome[i].allele_mutation_prob = m_prob;
	}

	return;
}



void genetic_pop_set_allele_min_double(organism_t *m_entity, const int m_num_chrom, const double m_min)
{

	if (!m_entity)
		die("Null pointer to entity structure passed.");

	for (int i = 0; i < m_num_chrom; i++)
	{
		m_entity->chromosome[i].allele_min_double = m_min;
	}
	return;
}

void genetic_organism_set_allele_max_double(organism_t *m_entity, const int m_num_chrom, const double m_max)
{

	if (!m_entity)
		die("Null pointer to entity structure passed.");

	for (int i = 0; i < m_num_chrom; i++)
	{
		m_entity->chromosome[i].allele_max_double = m_max;
	}
	return;
}

bool genetic_genocide(population_t *m_pop, size_t target_size)
{
	if (!m_pop)
		return false;

	elog(log_level_extreme, "The population is being culled from %zu to %zu organisms!", m_pop->size, target_size);

	/**
	 * Dereference the structures of the least fit population members until the desired population size is reached.
	 */
	while (m_pop->size > target_size)
	{
		genetic_free_rank(m_pop, m_pop->size - 1);
	}

	return true;
}


bool genetic_genocide_by_fitness(population_t *m_pop, double m_minfitness)
{
	if (!m_pop)
		return false;

	elog(log_level_extreme, "The population is being culled to minimum fitness %f!", m_minfitness);

	/**
	 * Dereference the structures relating to the least fit population members until
	 * the desired minimum fitness level is reached.
	 */
	while (m_pop->size > 0 && m_pop->sorted_organisms[m_pop->size - 1]->fitness < m_minfitness)
	{
		genetic_free_rank(m_pop, m_pop->size - 1);
	}

	elog(log_level_extreme, "After genocide, the population size is %zu", m_pop->size);

	return true;
}

bool genetic_close(population_t *m_pop)
{

	if (!m_pop)
		die("Null pointer to population structure passed.");

	elog(log_level_extreme, "This population is becoming extinct!");

	/*
	 * Dereference/free everything.
	 */
	if (!genetic_genocide(m_pop, 0))
	{
		elog(log_level_normal, "This population is already m_pop!");
	}
	else
	{
		e_free(m_pop->raw_organisms);
		e_free(m_pop->sorted_organisms);
		mem_chunk_destroy(m_pop->organism_memory);
		if (m_pop->chromo_chunk)
			mem_chunk_destroy(m_pop->chromo_chunk);
		if (m_pop->chromoarray_chunk)
			mem_chunk_destroy(m_pop->chromoarray_chunk);

		if (m_pop->de_params)
			e_free(m_pop->de_params);

		if (m_pop->data)
		{
			if (m_pop->population_data_destructor)
			{
				m_pop->population_data_destructor(m_pop->data);
			}
			else
			{
				elog(log_level_warning, "User data field is not empty. (Potential memory leak)");
			}
		}

		omp_destroy_lock(&m_pop->lock);
		omp_destroy_lock(&m_pop->chromo_chunk_lock);

		e_free(m_pop);
	}

	return true;
}


void ebex_genetic_init(void)
{

#pragma omp single
	{
		if (ebex_filter_openmp_initialised == false)
		{
			mem_chunk_init_openmp();
			ran_init();

			ebex_filter_openmp_initialised = true;
		}
	}

	return;
}

organism_t *ebex_genetic_get_mean_organism(population_t *m_pop)
{
	organism_t *adam;
	size_t original_size = m_pop->size;

	adam = genetic_get_new_organism(m_pop);
	for (size_t chromosome = 0; chromosome < m_pop->num_chromosomes; chromosome++)
	{
		for (size_t allele = 0; allele < adam->chromosome[chromosome].length; allele++)
		{
			switch(adam->chromosome[chromosome].type)
			{
				case CHROMOSOME_TYPE_DOUBLE:
					adam->chromosome[chromosome].dbl_allele[allele] = genetic_allele_mean_double(m_pop, chromosome, allele, original_size);
					break;
				case CHROMOSOME_TYPE_QUATERNION:
					genetic_allele_mean_quaternion(m_pop, chromosome, allele, &(adam->chromosome[chromosome].q_allele[allele]), original_size);
					break;
				default:
					die("Unknown chromosome type!");
					break;

			}
		}
	}

	return adam;
}

